'''
author:        wangchenyang <cy-wang21@mails.tsinghua.edu.cn>
date:          2024-08-06
Copyright © Department of Physics, Tsinghua University. All rights reserved
'''

''' Plot 1D GBZ as schematic diagram '''

import numpy as np
import matplotlib.pyplot as plt
import sys
import BerryPy.TightBinding as tb
sys.path.append('../../puncture-calculation')
import GBZ_manifold as Gm
import poly_tools as pt
import partial_GBZ_solver as pGs

def plot_GBZ():
    intra_cell = [
        [0, 1, 1],
        [1, 0, 1]
    ]
    inter_cell = [
        [0, 0, 0.5, (1,)],
        [0, 0, -0.5, (-1,)],
        [1, 1, 0.3, (1,)],
        [1, 1, 0.15, (-1,)],
        [0, 0, 0.15, (2,)],
        [0, 0, -0.15, (-2,)],
        [1, 1, 0.13, (2,)],
        [1, 1, 0.115, (-2,)]
    ]
    model = tb.TightBindingModel(1, 2, [[1]], intra_cell, inter_cell)

    coords, degs = model.get_characteristic_polynomial_data()
    char_poly = pt.CLaurent(2)
    char_poly.set_Laurent_by_terms(
        pt.CScalarVec(coords), 
        pt.CLaurentIndexVec(degs.flatten())
    )

    GBZ_segments, params, aGBZ_segments = pGs.solve_GBZ('GBZ',
                    char_poly)
    GBZ_segments, params, _ = pGs.expand_solution(GBZ_segments, aGBZ_segments, params) 
    all_segments = pGs.get_GBZ_segments(GBZ_segments, params, 0.1, pGs.aGBZEquation1DGBZCP1(char_poly))
    all_loops = pGs.GBZ_segments_to_loops(all_segments)
    all_loops = pGs.batch_glue(all_loops)

    for curr_loop in all_loops:
        all_points = np.zeros((len(curr_loop.point_vec),2), dtype=complex)
        for point_id, point in enumerate(curr_loop.point_vec):
            Gm.GBZ_to_chart(point, [0])
            all_points[point_id, :] = point.coords
        
        plt.plot(all_points[:,1].real, all_points[:,1].imag)
    
    plt.savefig("Figures/schematic-GBZ-1D.pdf")

if __name__ == '__main__':
    plot_GBZ()